package vectorstore

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"

	"github.com/cloudwego/eino/schema"
	"github.com/sirupsen/logrus"
)

// FaissIndexer Faiss索引器实现
type FaissIndexer struct {
	config      *FaissIndexerConfig
	documents   map[string]*schema.Document
	vectors     [][]float64
	vectorToID  map[int]string // 向量索引到文档ID的映射
	initialized bool
}

// NewFaissIndexer 创建新的Faiss索引器
func NewFaissIndexer(config *FaissIndexerConfig) (*FaissIndexer, error) {
	logrus.WithField("index_path", config.IndexPath).WithField("dimension", config.Dimension).Info("初始化Faiss索引器")

	// 确保目录存在
	if err := os.MkdirAll(filepath.Dir(config.IndexPath), 0755); err != nil {
		return nil, fmt.Errorf("创建索引目录失败: %w", err)
	}

	if err := os.MkdirAll(config.StoragePath, 0755); err != nil {
		return nil, fmt.Errorf("创建存储目录失败: %w", err)
	}

	return &FaissIndexer{
		config:     config,
		documents:  make(map[string]*schema.Document),
		vectorToID: make(map[int]string),
	}, nil
}

// AddDocuments 添加文档到索引
func (f *FaissIndexer) AddDocuments(ctx context.Context, docs []*schema.Document) error {
	logrus.WithField("count", len(docs)).Info("Faiss索引器添加文档")

	for _, doc := range docs {
		if doc.ID == "" {
			return fmt.Errorf("文档缺少ID")
		}
		if doc.Embedding == nil || len(doc.Embedding) != f.config.Dimension {
			return fmt.Errorf("文档嵌入向量维度不匹配，期望: %d, 实际: %d", f.config.Dimension, len(doc.Embedding))
		}

		// 存储文档
		f.documents[doc.ID] = doc
		// 存储向量
		vecIndex := len(f.vectors)
		f.vectors = append(f.vectors, doc.Embedding)
		f.vectorToID[vecIndex] = doc.ID
	}

	f.initialized = true
	return nil
}

// UpdateDocuments 更新文档
func (f *FaissIndexer) UpdateDocuments(ctx context.Context, docs []*schema.Document) error {
	logrus.WithField("count", len(docs)).Info("Faiss索引器更新文档")

	// 先删除旧文档
	idsToDelete := make([]string, len(docs))
	for i, doc := range docs {
		idsToDelete[i] = doc.ID
	}

	if err := f.DeleteDocuments(ctx, idsToDelete); err != nil {
		return err
	}

	// 再添加新文档
	return f.AddDocuments(ctx, docs)
}

// DeleteDocuments 删除文档
func (f *FaissIndexer) DeleteDocuments(ctx context.Context, ids []string) error {
	logrus.WithField("count", len(ids)).Info("Faiss索引器删除文档")

	// 标记要删除的文档
	deletedIDs := make(map[string]bool)
	for _, id := range ids {
		deletedIDs[id] = true
		delete(f.documents, id)
	}

	// 重建向量和映射
	newVectors := [][]float64{}
	newVectorToID := make(map[int]string)

	for i, id := range f.vectorToID {
		if !deletedIDs[id] {
			newVecIndex := len(newVectors)
			newVectors = append(newVectors, f.vectors[i])
			newVectorToID[newVecIndex] = id
		}
	}

	f.vectors = newVectors
	f.vectorToID = newVectorToID

	return nil
}

// Search 搜索相似文档（简单的暴力搜索实现）
func (f *FaissIndexer) Search(ctx context.Context, queryVector []float64, topK int) ([]*schema.Document, error) {
	logrus.WithField("topK", topK).Info("Faiss索引器搜索文档")

	if !f.initialized {
		return []*schema.Document{}, nil
	}

	if len(queryVector) != f.config.Dimension {
		return nil, fmt.Errorf("查询向量维度不匹配，期望: %d, 实际: %d", f.config.Dimension, len(queryVector))
	}

	// 简单的余弦相似度搜索
	type scoredResult struct {
		index int
		score float64
	}

	var results []scoredResult

	for i, vector := range f.vectors {
		score := cosineSimilarity(queryVector, vector)
		results = append(results, scoredResult{index: i, score: score})
	}

	// 按分数排序
	sort.Slice(results, func(i, j int) bool {
		return results[i].score > results[j].score
	})

	// 返回前topK个结果
	resultDocs := []*schema.Document{}
	count := 0

	for _, result := range results {
		if count >= topK {
			break
		}

		if docID, exists := f.vectorToID[result.index]; exists {
			if doc, exists := f.documents[docID]; exists {
				resultDocs = append(resultDocs, doc)
				count++
			}
		}
	}

	return resultDocs, nil
}

// GetDocuments 根据ID获取文档
func (f *FaissIndexer) GetDocuments(ctx context.Context, ids []string) ([]*schema.Document, error) {
	logrus.WithField("count", len(ids)).Info("Faiss索引器获取文档")

	var results []*schema.Document
	for _, id := range ids {
		if doc, exists := f.documents[id]; exists {
			results = append(results, doc)
		}
	}

	return results, nil
}

// Save 保存索引
func (f *FaissIndexer) Save(ctx context.Context) error {
	logrus.WithField("path", f.config.IndexPath).Info("保存Faiss索引")

	// 保存向量数据
	vectorData := map[string]interface{}{
		"vectors":    f.vectors,
		"vectorToID": f.vectorToID,
		"dimension":  f.config.Dimension,
	}

	vectorDataBytes, err := json.Marshal(vectorData)
	if err != nil {
		return fmt.Errorf("序列化向量数据失败: %w", err)
	}

	if err := os.WriteFile(f.config.IndexPath, vectorDataBytes, 0644); err != nil {
		return fmt.Errorf("写入索引文件失败: %w", err)
	}

	// 保存文档数据
	docDataBytes, err := json.Marshal(f.documents)
	if err != nil {
		return fmt.Errorf("序列化文档数据失败: %w", err)
	}

	docFilePath := filepath.Join(f.config.StoragePath, "documents.json")
	if err := os.WriteFile(docFilePath, docDataBytes, 0644); err != nil {
		return fmt.Errorf("写入文档文件失败: %w", err)
	}

	return nil
}

// Load 加载索引
func (f *FaissIndexer) Load(ctx context.Context) error {
	logrus.WithField("path", f.config.IndexPath).Info("加载Faiss索引")

	// 检查索引文件是否存在
	if _, err := os.Stat(f.config.IndexPath); os.IsNotExist(err) {
		logrus.Info("索引文件不存在，使用空索引")
		return nil
	}

	// 读取向量数据
	vectorDataBytes, err := os.ReadFile(f.config.IndexPath)
	if err != nil {
		return fmt.Errorf("读取索引文件失败: %w", err)
	}

	var vectorData map[string]interface{}
	if err := json.Unmarshal(vectorDataBytes, &vectorData); err != nil {
		return fmt.Errorf("解析索引数据失败: %w", err)
	}

	// 读取文档数据
	docFilePath := filepath.Join(f.config.StoragePath, "documents.json")
	docDataBytes, err := os.ReadFile(docFilePath)
	if err != nil {
		return fmt.Errorf("读取文档文件失败: %w", err)
	}

	if err := json.Unmarshal(docDataBytes, &f.documents); err != nil {
		return fmt.Errorf("解析文档数据失败: %w", err)
	}

	f.initialized = true
	return nil
}

// Close 关闭索引器
func (f *FaissIndexer) Close() error {
	logrus.Info("关闭Faiss索引器")
	// 在这个实现中不需要特殊的关闭操作
	return nil
}
